﻿#include "algo/algo_runge_kutta.h"
#include "math/Polynomial.h"
#include "math/nonlinear_math.h"

CAGD::ClassicalRK::ClassicalRK(double k) : t_(0), k_(k), A_(4, 4), b_(4), c_(4)
{
    // RK 节点和权
    c_ = {0, 0.5, 0.5, 1};
    b_ = {1.0 / 6, 1.0 / 3, 1.0 / 3, 1.0 / 6};

    A_(1, 0) = 0.5;
    A_(2, 1) = 0.5;
    A_(3, 2) = 1;
}

CAGD::GTimer CAGD::ClassicalRK::shift(FuncXXt f, GTimer U)
{
    // 计算每个 y
    Vector y1 = f(*U, t_);
    Vector y2 = f(*U + y1 * A_(1, 0) * k_, t_ + c_[1] * k_);
    Vector y3 = f(*U + y2 * A_(2, 1) * k_, t_ + c_[2] * k_);
    Vector y4 = f(*U + y3 * A_(3, 2) * k_, t_ + c_[3] * k_);

    // 加权求和
    *(U + 1) = *U + (y1 * b_[0] + y2 * b_[1] + y3 * b_[2] + y4 * b_[3]) * k_;

    // 时间推进
    t_ += k_;

    return U + 1;
}

CAGD::ESDIRK::ESDIRK(double k) : s_(6), t_(0), k_(k), A_(6, 6), b_(6), c_(6)
{
    // RK 节点和权
    c_ = {0, 0.5, 83.0 / 250, 31.0 / 50, 17.0 / 20, 1};
    b_ = {82889.0 / 524892, 0, 15625.0 / 83664, 69875.0 / 102672, -2260.0 / 8211, 0.25};

    A_(1, 0) = 0.25;
    A_(1, 1) = 0.25;

    A_(2, 0) = 8611.0 / 62500;
    A_(2, 1) = -1743.0 / 31250;
    A_(2, 2) = 0.25;

    A_(3, 0) = 5012029.0 / 34652500;
    A_(3, 1) = -654441.0 / 2922500;
    A_(3, 2) = 174375.0 / 388108;
    A_(3, 3) = 0.25;

    A_(4, 0) = 15267082809.0 / 155376265600;
    A_(4, 1) = -71443401.0 / 120774400;
    A_(4, 2) = 730878875.0 / 902184768;
    A_(4, 3) = 2285395.0 / 8070912;
    A_(4, 4) = 0.25;

    for (int i = 0; i < 6; i++)
        A_(5, i) = b_[i];
}

CAGD::GTimer CAGD::ESDIRK::shift(FuncXXt f, GTimer U)
{
    // 存放 y 值，每个都初始化为 0
    std::vector<Vector> Y(s_, f(*U, t_));

    // 牛顿迭代器
    NewtonSolver fS;

    // 对角线循环
    for (int i = 0; i < s_; i++)
    {
        // 构造隐式方程
        FuncXX F = [=](Vector X) -> Vector {
            Vector ty = *U;
            for (int j = 0; j <= i; j++)
                ty = ty + Y[j] * A_(i, j) * k_;
            return X - f(ty, t_ + c_[i] * k_);
        };

        // 迭代求根
        fS(F, Y[i]);
    }

    // 加权求和
    *(U + 1) = *U;
    for (int i = 0; i < s_; i++)
        *(U + 1) = *(U + 1) + Y[i] * b_[i] * k_;

    // 时间推进
    t_ += k_;

    return U + 1;
}

CAGD::GaussLegendreRK::GaussLegendreRK(int s, double k) : s_(s), t_(0), k_(k), A_(s, s), b_(s), c_(s)
{
    // 获得首一 Gauss-Legendre 多项式
    Polynomial GL = CAGD::monic_legendre(s);

    // 使用牛顿法迭代求根
    NewtonSolver fS;

    // 依次计算 s 个根，从 0 开始迭代，通常会得到从小到大的根
    for (int i = 0; i < s; i++)
    {
        // 构造函数
        Func11 F = [=](double x) -> double { return GL.Value(x); };

        // 迭代求根
        c_[i] = fS(F, 0);

        // 除以因式 x - c[i] 进行降阶
        GL = GL / Polynomial({-c_[i], 1});
    }

    // 保险起见按照升序排列
    c_.Sort();

    // 初始化权
    for (int i = 0; i < s; i++)
        b_[i] = CAGD::lagrange(c_, i).Intergral(0, 1);

    // 初始化系数矩阵
    for (int i = 0; i < s; i++)
        for (int j = 0; j < s; j++)
            A_(i, j) = CAGD::lagrange(c_, j).Intergral(0, c_[i]);
}

CAGD::GTimer CAGD::GaussLegendreRK::shift(FuncXXt f, GTimer U)
{
    int N = U->Size();

    // 变量规模：共 s_ 个方程，每个方程是一个 N 元向量
    // Y = (y1 y2 ... ys)
    FuncXX F = [=](Vector Y) -> Vector {
        // 结果容器
        Vector res(Y.Size());

        // i ：第 i 个方程
        for (int i = 0; i < s_; i++)
        {
            // 获得初始向量
            Vector yi = *U;

            for (int j = 0; j < s_; j++)
                for (int n = 0; n < N; n++)
                    yi[n] += k_ * A_(i, j) * Y[j * N + n];

            // 计算 yi
            yi = f(yi, t_ + k_ * c_[i]);

            // 得到第 i 方程
            for (int n = 0; n < N; n++)
                res[i * N + n] = Y[i * N + n] - yi[n];
        }
        return res;
    };

    // 使用 f(U^n,t) 作为初值
    Vector Y0 = f(*U, t_);
    Vector Y(s_ * N);

    for (int i = 0; i < s_; i++)
        for (int n = 0; n < N; n++)
            Y[i * N + n] = Y0[n];

    // 牛顿法迭代求根
    NewtonSolver fS;
    fS(F, Y);

    // 加权推进
    *(U + 1) = *U;
    for (int i = 0; i < s_; i++)
        for (int n = 0; n < N; n++)
            (*(U + 1))[n] += k_ * b_[i] * Y[i * N + n];

    // 时间推进
    t_ += k_;

    return U + 1;
}